package com.psmz.iwi.docanalyzer.ui;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreeNode;

import com.psmz.iwi.docanalyzer.core.document.entities.Document;
import com.psmz.iwi.docanalyzer.core.document.entities.NumberOfWordsClassificationGroup;

public class ClassifiedDocumentsTreeModel extends DefaultTreeModel {

	private static final String ROOT_NODE_OBJECT = "Repozytorium";
	private List<NumberOfWordsClassificationGroup> groups = new ArrayList<NumberOfWordsClassificationGroup>();

	public ClassifiedDocumentsTreeModel(
			List<NumberOfWordsClassificationGroup> groups) {
		super(new DefaultMutableTreeNode(ROOT_NODE_OBJECT));
		this.groups = new ArrayList<NumberOfWordsClassificationGroup>(groups);

	}

	public ClassifiedDocumentsTreeModel(TreeNode root) {
		super(root);
	}

	@Override
	public Object getRoot() {
		return ROOT_NODE_OBJECT;
	}

	@Override
	public Object getChild(Object parent, int index) {
		if (parent instanceof String) {
			return groups.get(index);
		} else if (parent instanceof NumberOfWordsClassificationGroup) {
			NumberOfWordsClassificationGroup group = (NumberOfWordsClassificationGroup) parent;
			return group.getDocumentOnPosition(index);
		} else {
			return null;
		}
	}

	@Override
	public int getChildCount(Object parent) {
		if (parent instanceof String) {
			return groups.size();
		} else if (parent instanceof NumberOfWordsClassificationGroup) {
			NumberOfWordsClassificationGroup group = (NumberOfWordsClassificationGroup) parent;
			return group.getNumberOfDocuments();
		} else {
			return 0;
		}
	}

	@Override
	public boolean isLeaf(Object node) {
		if (node instanceof Document) {
			return true;
		} else {
			return false;
		}
	}

	@Override
	public int getIndexOfChild(Object parent, Object child) {
		if (parent instanceof String
				&& child instanceof NumberOfWordsClassificationGroup) {
			return groups.indexOf(child);
		} else if (parent instanceof NumberOfWordsClassificationGroup
				&& child instanceof Document) {
			NumberOfWordsClassificationGroup group = (NumberOfWordsClassificationGroup) parent;
			return group.getDocumentIndex((Document) child);
		}
		return -1;
	}

	public void reloadData(Collection<NumberOfWordsClassificationGroup> groups) {
		this.groups = new ArrayList<NumberOfWordsClassificationGroup>(groups);
		reload();
	}

}
